from io import BytesIO
import threading
from tkinter import Button, Tk, Label, messagebox, filedialog
import windnd
import zipfile
import os
import traceback
import subprocess

mainWindow = None


def checkFFmpegEnv():
    try:
        p1 = subprocess.Popen(["ffmpeg", "-version"])
    except:
        try:
            p2 = subprocess.Popen(["avconv", "-version"])
        except:
            return False
        else:
            p2.wait()
            return p2.returncode == 0
    else:
        p1.wait()
        return p1.returncode == 0


def checkPip():
    try:
        p = subprocess.Popen("pip --version")
    except:
        return False
    else:
        p.wait()
        return p.returncode == 0


def installDepByName(name: str):
    try:
        return os.system("pip install -i https://pypi.tuna.tsinghua.edu.cn/simple "+name) == 0
    except:
        return False


class CheckEnvTk:
    # 检查环境窗口类
    def __init__(self):
        global mainWindow
        mainWindow.withdraw()

        tk = Tk()
        tk.title("检测环境窗口")
        tk.geometry("400x300")
        tk.protocol("WM_DELETE_WINDOW", self.onClosed)
        self.tk = tk

        label0 = Label(tk,text="以下每项都要安装，程序才能正常工作\n")
        label0.pack()

        label1 = Label(tk, text="依赖库相关")
        label1.pack()
        installDepBtn = Button(tk, text="安装依赖库", command=self.installDep)
        installDepBtn.pack()
        self.installDepBtn = installDepBtn

        label2 = Label(tk, text="ffmpeg相关")
        label2.pack()
        Label(tk,text="请先检测，不通过再安装").pack()
        ffBtn = Button(tk, text="ffmpeg检测", bg="yellow",
                       command=self.checkFFmpegEnvForBtn)
        ffBtn.pack()
        self.ffBtn = ffBtn
        installffBtn = Button(tk, text="自动安装ffmpeg",
                              command=self.installFFmpeg)
        installffBtn.pack()

        label3 = Label(tk, text="数据集相关")
        label3.pack()
        pthBtn = Button(tk, text="自动下载数据集", command=self.downloadPth)
        pthBtn.pack()

    def onClosed(self):
        global mainWindow
        mainWindow.wm_deiconify()
        self.tk.destroy()

    def installDep(self):
        ask = messagebox.askokcancel("安装依赖库", "是否安装依赖库（只需要安装一次）")
        if not ask:
            return
        if not checkPip():
            messagebox.showerror(
                "pip不可用", "您的pip不可用，请检查python环境是否正确安装，建议使用搜索引擎查询此问题")
            return

        def install():
            deps = ["h5py", "pandas", "librosa", "numba", "mido",
                "mir_eval", "matplotlib", "torchlibrosa", "sox", "torch"]
            for dep in deps:
                self.installDepBtn.config(text="正在安装 %s"%dep)
                if not installDepByName(dep):
                    messagebox.showerror("安装错误", "依赖库"+dep+"安装错误，请检查错误或自行安装")
                    return
            self.installDepBtn.config(text="安装依赖库")
            messagebox.showinfo("安装完成", "所有依赖库均已安装成功")
        
        t = threading.Thread(target=install)
        t.setDaemon(True)
        t.start()

    def checkFFmpegEnvForBtn(self):
        self.ffBtn.config(text="ffmpeg检测中")
        if checkFFmpegEnv():
            self.ffBtn.config(bg="green", text="ffmpeg检测通过")
        else:
            self.ffBtn.config(bg="red", text="ffmpeg检测失败")
            messagebox.showinfo(
                "可以自动安装ffmpeg", "点击下方按钮可以帮您自动安装ffmpeg，若您想手动安装请无视此窗口")

    def installFFmpeg(self):
        ask = messagebox.askokcancel(
            "安装ffmpeg", "是否自动下载安装ffmpeg，此方法会将ffmpeg放置到“C:/ffmpeg”中，若不接受请取消（下载中程序会处于无响应状态）")
        if not ask:
            return
        if not installDepByName("requests"):
            messagebox.showerror("错误", "无法安装requests，安装中止")
            return
        # 下载ffmpeg压缩包
        import requests
        url = "https://hub.fastgit.org/BtbN/FFmpeg-Builds/releases/download/autobuild-2021-08-21-12-22/ffmpeg-N-103313-ga1094eaa30-win64-lgpl-shared.zip"
        req = requests.get(url, stream=True)
        data = BytesIO()
        # 写入内存
        for chunk in req.iter_content(1048576):
            if chunk:
                data.write(chunk)
        # 检查是否需要创建文件夹
        if not os.path.exists("C:\\ffmpeg"):
            os.mkdir("C:\\ffmpeg")
        # 读取zip内容
        zip = zipfile.ZipFile(data)
        names = ["avcodec-59.dll", "avdevice-59.dll", "avfilter-8.dll", "avformat-59.dll",
                 "avutil-57.dll", "ffmpeg.exe", "ffplay.exe", "ffprobe.exe", "swresample-4.dll", "swscale-6.dll"]
        for name in names:
            # 将所需文件提取到目标目录
            file = zip.read(
                "ffmpeg-N-103313-ga1094eaa30-win64-lgpl-shared/bin/"+name)
            with open("C:\\ffmpeg\\"+name, "wb") as f:
                f.write(file)
        data.close()
        # 获取环境变量
        var = str(subprocess.check_output(
            ["powershell", "(Get-Item HKCU:\Environment).GetValue(\"Path\",$null,\"DoNotExpandEnvironmentNames\")"]), "utf-8").strip("\r\n")
        if var.endswith(";"):
            var += "C:\\ffmpeg;"
        else:
            var += ";C:\\fmpeg;"
        # 写入path变量值
        p = subprocess.Popen(["powershell", "setx Path \""+var+"\""])
        p.wait()
        if p.returncode != 0:
            messagebox.showerror("安装失败", "环境变量配置出现错误，请尝试手动安装")
            return
        messagebox.showinfo("安装成功", "ffmpeg已经安装到C:\\ffmpeg，请重启本程序以使安装生效")

    def downloadPth(self):
        if not messagebox.askokcancel("提示","是否下载数据集，文件大小为165MB左右（下载中窗口无响应）"):
            return
        if not installDepByName("requests"):
            messagebox.showerror("错误", "无法安装requests，安装中止")
            return
        import requests
        checkpoint_path='{}/piano_transcription_inference_data/note_F1=0.9677_pedal_F1=0.9186.pth'.format(os.getenv("USERPROFILE"))
        if not os.path.exists(checkpoint_path) or os.path.getsize(checkpoint_path) < 1.6e8:
            if not os.path.exists(os.path.dirname(checkpoint_path)):
                os.mkdir(os.path.dirname(checkpoint_path))
            zenodo_path = 'https://zenodo.org/record/4034264/files/CRNN_note_F1%3D0.9677_pedal_F1%3D0.9186.pth?download=1'
            req = requests.get(zenodo_path,stream=True)
            file = open(checkpoint_path,"wb")
            for chunk in req.iter_content(1048576):
                if(chunk):
                    file.write(chunk)
            file.close()
        messagebox.showinfo("下载完成","下载已完成，存放到%s"%checkpoint_path)


    def mainloop(self):
        self.tk.mainloop()


def openCheckEnvTk():
    tk = CheckEnvTk()
    tk.mainloop()


class MainTk:
    # 主窗口类
    def __init__(self):
        tk = Tk()
        tk.title("AI扒谱辅助工具")
        tk.geometry("400x300")
        windnd.hook_dropfiles(tk, self.dropFile)
        self.tk = tk

        self.running = False

        label1 = Label(tk, text="第一次运行请检查环境")
        label1.pack()

        checkBtn = Button(tk, text="检查环境", command=openCheckEnvTk)
        checkBtn.pack()

        label2 = Label(tk, text="AI扒谱")
        label2.pack()
        label3 = Label(tk, text="可以通过下方按钮选择文件，或直接将文件拖入此窗口")
        label3.pack()
        pickBtn = Button(tk, text="选择文件", command=self.pickFile)
        pickBtn.pack()

        statusLabel = Label(tk, text="状态：未运行")
        statusLabel.pack()
        self.statusLabel = statusLabel

    def pickFile(self):
        if self.running:
            return
        file = filedialog.askopenfilename()
        if file.strip() == "":
            return
        self.handleFile(file)

    def dropFile(self, files):
        if self.running:
            return
        file = files[0]
        try:
            file = str(file, "gbk")
        except:
            file = str(file, "utf-8")
        self.handleFile(file)

    def handleFile(self, file: str):
        ask = messagebox.askokcancel(
            "提示", "你选择了文件"+file+"，是否确认并开始运行？")
        if not ask:
            return
        if not os.path.exists(file):
            messagebox.showerror("错误", "文件不存在")
            return
        t = threading.Thread(target=self.start, args=[file])
        t.setDaemon(True)
        t.start()
        self.running = True
        self.updateStatusLabel("状态：正在开始")

    def start(self, file: str):
        try:
            import audioread
            from piano_transcription_inference_modified import PianoTranscription, sample_rate, load_audio
            piano = PianoTranscription(device="cpu")
            (audio, _) = load_audio(file, sr=sample_rate)
            out = os.path.splitext(file)[0]+".mid"
            piano.transcribe(audio, out, self.progress)
            messagebox.showinfo("提示", "成功，midi文件输出到"+out)
        except audioread.exceptions.NoBackendError:
            messagebox.showerror("错误","ffmpeg未安装，或您选择了非音频文件，请检查后重新运行")
        except BaseException as ex:
            messagebox.showerror("错误",traceback.format_exc(ex))
        finally:
            self.running = False
            self.updateStatusLabel("状态：未运行")

    def progress(self, a, b):
        self.updateStatusLabel("状态：运行中 %d / %d" % (a, b))

    def updateStatusLabel(self, text):
        self.statusLabel.config(text=text)

    def mainloop(self):
        self.tk.mainloop()

    def withdraw(self):
        self.tk.withdraw()

    def wm_deiconify(self):
        self.tk.wm_deiconify()


def main():
    # 运行主窗口
    global mainWindow
    mainWindow = MainTk()
    mainWindow.mainloop()


if __name__ == "__main__":
    main()
